package com.flink.cdc;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.streaming.connectors.kafka.KafkaContextAware;
import org.apache.flink.streaming.connectors.kafka.KafkaSerializationSchema;
import org.apache.kafka.clients.producer.ProducerRecord;

import java.nio.charset.StandardCharsets;

import javax.annotation.Nullable;

public class SimpleStringSerializationSchemaWrapper
        implements KafkaSerializationSchema<String>, KafkaContextAware<String> {

    private final SerializationSchema<String> serializationSchema = new SimpleStringSchema();
    private boolean writeTimestamp;

    public SimpleStringSerializationSchemaWrapper(boolean writeTimestamp) {
        this.writeTimestamp = writeTimestamp;
    }

    @Override
    public void open(SerializationSchema.InitializationContext context) throws Exception {
        serializationSchema.open(context);
    }

    @Override
    public ProducerRecord<byte[], byte[]> serialize(String element, @Nullable Long timestamp) {
        byte[] serialized = serializationSchema.serialize(element);

        final Long timestampToWrite;
        if (writeTimestamp) {
            timestampToWrite = timestamp;
        } else {
            timestampToWrite = null;
        }

        String targetTopic = getTargetTopic(element);
        // 按topic(db.table)名来分区
        byte[] key = targetTopic.getBytes(StandardCharsets.UTF_8);
        return new ProducerRecord<>(targetTopic, null, timestampToWrite, key, serialized);
    }

    @Override
    public String getTargetTopic(String element) {
        JSONObject payload = JSON.parseObject(element);
        JSONObject source = payload.getJSONObject("source");
        String db = source.getString("db");
        String table = source.getString("table");
        return db + "." + table;
    }

    public void setWriteTimestamp(boolean writeTimestamp) {
        this.writeTimestamp = writeTimestamp;
    }
}
